"""
改变via的标注文件，将metadata下的键都统一改成image<视频名>_<对象序号>的形式
"""
import os
import json


def modify_metadata_key(metadata):
  # 创建一个字典来记录已修改的键
  key_count = {}
  new_metadata = {}

  for old_key, value in metadata.items():
    # 提取原始键名中的数字部分
    parts = old_key.split('_')

    if len(parts) == 2 and parts[0].startswith('image') and parts[1].isdigit():
      # 如果键符合image<数字>_<数字>格式，则保持不变
      new_key = old_key
    else:
      # 否则，提取数字部分，创建新的键名
      try:
        prefix_num = int(parts[0]) if parts[0].isdigit() else None
      except ValueError:
        prefix_num = None

      if prefix_num is not None:
        # 尝试生成新的格式 image<数字>_<数字>
        new_key = f"image{prefix_num}_{1}"
        # 处理键已存在的情况
        while new_key in new_metadata:
          count = key_count.get(new_key, 0) + 1
          new_key = f"image{prefix_num}_{count}"
          key_count[new_key] = count

    # 将修改后的键值对加入新的字典
    new_metadata[new_key] = value

  return new_metadata


def process_json_file(file_path, new_file_path):
  with open(file_path, 'r') as file:
    data = json.load(file)

  # 修改metadata部分
  if 'metadata' in data:
    data['metadata'] = modify_metadata_key(data['metadata'])

  # 保存修改后的数据到新的文件
  with open(new_file_path, 'w') as file:
    json.dump(data, file, indent=4)


def process_directory(root_dir):
  # 遍历根目录下的所有子文件夹
  for subfolder_name in os.listdir(root_dir):
    subfolder_path = os.path.join(root_dir, subfolder_name)
    if os.path.isdir(subfolder_path):
      # 查找子文件夹下的_finish.json文件
      json_file_name = f"{subfolder_name}_finish.json"
      json_file_path = os.path.join(subfolder_path, json_file_name)

      if os.path.exists(json_file_path):
        # 创建新的文件路径 <子文件夹名>_finish_correct.json
        new_json_file_name = f"{subfolder_name}_finish_correct.json"
        new_json_file_path = os.path.join(subfolder_path, new_json_file_name)

        # 处理并保存到新的文件
        process_json_file(json_file_path, new_json_file_path)


if __name__ == "__main__":
  # 请将以下路径替换为实际的文件夹路径
  root_directory = r"D:\file\postgrad\experiment\ava_dataset\try"
  process_directory(root_directory)
